#ifndef THREADGROUP_H_
#define THREADGROUP_H_

#include <thread>
#include <memory>
#include <vector>
#include <mutex>
#include <functional>
#include <algorithm>
#include <unordered_map>

namespace toolkit {

/**
 * @brief 线程组类，用于管理一组线程
 */
class ThreadGroup {
public:
    /**
     * @brief 构造函数
     */
    ThreadGroup() = default;

    /**
     * @brief 析构函数，等待所有线程退出
     */
    ~ThreadGroup() {
        joinAll();
    }

    /**
     * @brief 创建线程并加入线程组
     * @param func 线程函数
     * @param args 线程函数参数
     * @return 创建的线程ID
     */
    template<typename Function, typename... Args>
    std::thread::id createThread(Function&& func, Args&&... args) {
        std::unique_lock<std::mutex> lock(_mutex);
        
        auto thread_ptr = std::make_shared<std::thread>(
            std::forward<Function>(func), 
            std::forward<Args>(args)...
        );
        
        std::thread::id thread_id = thread_ptr->get_id();
        _threads[thread_id] = thread_ptr;
        
        return thread_id;
    }

    /**
     * @brief 从线程组中移除指定线程
     * @param thread_id 要移除的线程ID
     * @param wait 是否等待线程结束
     * @return 是否成功移除
     */
    bool removeThread(std::thread::id thread_id, bool wait = true) {
        std::shared_ptr<std::thread> thread_ptr;
        
        {
            std::unique_lock<std::mutex> lock(_mutex);
            auto it = _threads.find(thread_id);
            if (it == _threads.end()) {
                return false;
            }
            
            thread_ptr = it->second;
            _threads.erase(it);
        }
        
        if (wait && thread_ptr->joinable()) {
            thread_ptr->join();
        }
        
        return true;
    }

    /**
     * @brief 等待所有线程退出
     */
    void joinAll() {
        std::unordered_map<std::thread::id, std::shared_ptr<std::thread>> threads_copy;
        
        {
            std::unique_lock<std::mutex> lock(_mutex);
            threads_copy = _threads;
            _threads.clear();
        }
        
        for (auto& pair : threads_copy) {
            if (pair.second->joinable()) {
                pair.second->join();
            }
        }
    }

    /**
     * @brief 判断指定线程是否在线程组中
     * @param thread_id 线程ID
     * @return 是否在线程组中
     */
    bool hasThread(std::thread::id thread_id) const {
        std::unique_lock<std::mutex> lock(_mutex);
        return _threads.find(thread_id) != _threads.end();
    }

    /**
     * @brief 判断当前线程是否在线程组中
     * @return 是否在线程组中
     */
    bool isCurrentThreadIn() const {
        return hasThread(std::this_thread::get_id());
    }

    /**
     * @brief 获取线程组中线程数量
     * @return 线程数量
     */
    size_t size() const {
        std::unique_lock<std::mutex> lock(_mutex);
        return _threads.size();
    }

private:
    mutable std::mutex _mutex;
    std::unordered_map<std::thread::id, std::shared_ptr<std::thread>> _threads;
};

} /* namespace toolkit */

#endif /* THREADGROUP_H_ */